package com.uwlaser.config.sharding;

import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;
import org.springframework.stereotype.Component;

import java.util.*;

/**
 * @author zk-Qin
 * @create 2022-06-11 18:56
 */
@Component
public class DataComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm {
    /**
     * 产品类型的分片
     */
    private static Map<String, String> PRODUCT_TYPE_MAP = new HashMap<String, String>() {
        {
            put("P","p");
            put("M", "m");
            put("C", "c");
        }
    };

    /**
     * 工段类型的分片
     */
    private static Map<String,String> SECTION_CODE_MAP = new HashMap<String, String>() {
        {
            put("PACK","p");
            put("MODULE", "m");
            put("CELL", "c");
            put("cell-section", "c");
            put("module-section", "m");
            put("pack-section", "p");
        }
    };

    @Override
    public Collection<String> doSharding(Collection collection, ComplexKeysShardingValue complexKeysShardingValue) {

        List<String> result = new ArrayList<>();
        //获取分表列 分片表规则
        List<String> productTypes = this.getShardingValue(complexKeysShardingValue, "product_type");
        List<String> sectionCodes = this.getShardingValue(complexKeysShardingValue, "section_code");

        collection.forEach(table -> {
            if (productTypes.isEmpty() && sectionCodes.isEmpty()) {
                //查询SQL 一个分片键都没有情况
                result.add(String.valueOf(table));
            } else if(productTypes.isEmpty()) {
                //查询SQL 只有工段 一个分片建
                String tableSuffix = "_"+ this.SECTION_CODE_MAP.get(sectionCodes.get(0));
                String tableTemp = String.valueOf(table);
                if(tableTemp.endsWith(tableSuffix)){
                    result.add(tableTemp);
                }
            } else if (sectionCodes.isEmpty()){
                //查询SQL 只有user_id 一个分片建
                String tableSuffix = "_"+this.PRODUCT_TYPE_MAP.get(productTypes.get(0));
                String tableTemp = String.valueOf(table);
                if(tableTemp.endsWith(tableSuffix)){
                    result.add(tableTemp);
                }
            } else {
                //全部都有分片键存在
                String tableSuffix = "_"+this.PRODUCT_TYPE_MAP.get(productTypes.get(0));
                String tableTemp = String.valueOf(table);
                if(tableTemp.endsWith(tableSuffix)){
                    result.add(tableTemp);
                }
            }
        });


        return result;
    }

    //获取分片键的值
    private List<String> getShardingValue(ComplexKeysShardingValue complexKeysShardingValue, String key) {
        List<String> result = new ArrayList<>();
        if (complexKeysShardingValue.getColumnNameAndShardingValuesMap().size() > 0 && complexKeysShardingValue.getColumnNameAndShardingValuesMap().containsKey(key)) {
            Collection<String> value = (Collection<String>) complexKeysShardingValue.getColumnNameAndShardingValuesMap().get(key);
            ArrayList<String> strings = new ArrayList<>(value);
            for (String v : value) {
                result.add(v);
            }
            return result;
        }
        return result;
    }
}
